"use strict";

import { app } from "../../../scripts/app.js";
import { api } from "../../../scripts/api.js";
import { loadMetadata } from "./libs/workflow-editor.js";
import JSON5 from "./libs/json5.min.js";

const PARSER_TYPES = [
  "LoadBooleanFromImage",
  "LoadIntFromImage",
  "LoadFloatFromImage",
  "LoadStringFromImage",
  "LoadComboFromImage",
];

const IMAGE_LOADER_TYPES = [
  "LoadImage",
  "LoadImageMask",
  "LoadImage //Inspire",
  "Load image with metadata [Crystools]",
  "Image Load", // was-node-suite-comfyui
  "PutImage",
];

function isSupportedNode(node) {
  if (!node || !node.comfyClass) {
    return false;
  }
  if (IMAGE_LOADER_TYPES.indexOf(node.comfyClass) > -1) {
    return true;
  }
  return false;
}

function initParser() {
  try {
    if (this.statics && this.statics.isInitialized) {
      return;
    }

    const self = this;

    this.statics = {
      isInitialized: false,
      imageNode: undefined,
    };

    this.statics.getImageNode = (function() {
      let input = this.inputs?.find(e => e.name === "image");
      if (!input || !input.link) {
        return;
      }
      let linkId = input.link;
      let link = app.graph.links.get(linkId);
      let targetId = link.origin_id;
      let target = app.graph._nodes.find(e => e.id === targetId);

      return isSupportedNode(target) ? target : undefined;

      // deprecated
      // while(target && IMAGE_LOADER_TYPES.indexOf(target.comfyClass) === -1) {
      //   input = target.inputs?.find(e => ["pixels", "image"].indexOf(e.name) > -1);
      //   if (!input || !input.link) {
      //     return;
      //   }
      //   linkId = input.link;
      //   link = graph.links.find(e => e && e.id === linkId);
      //   targetId = link.origin_id;
      //   target = app.graph._nodes.find(e => e.id === targetId);
      // }
      // return target;
    }).bind(this);

    this.statics.setImageNode = (function(imageNode) {
      this.statics.imageNode = imageNode || this.statics.getImageNode();
      return this.statics.imageNode;
    }).bind(this);

    this.statics.getMetadata = (async function() {
      const node = this.statics.setImageNode();
      if (!node) {
        return;
      }
      if (node.parsedData) {
        return node.parsedData;
      }

      const filePath = getFilePath(node);
      if (!filePath) {
        return;
      }
      
      try {
        const { 
          fullname,
          dirname,
          filename,
          extname,
          width, 
          height, 
          info, 
          format, // "PNG"
        } = await loadMetadata(filePath);

        const prompt = JSON5.parse(info.prompt);
        const workflow = JSON5.parse(info.workflow);

        let notes = [];

        try {
          notes = workflow.nodes?.filter((e) => {
            return e.type === "Note";
          }).map((node) => {
            const id = node.id;
            const type = node.type;
            const values = {
              "text": node.widgets_values[0],
            }
            let title = node.title;
            if (!title && LiteGraph.registered_node_types[type]) {
              const n = LiteGraph.registered_node_types[type];
              title = n.title;
            }
            return {
              id,
              title,
              type,
              values,
            }
          }) || [];
        } catch(err) {
          console.error(err);
        }
        
        const nodes = Object.entries(prompt).map(e => {
          const id = parseInt(e[0]);
          const type = e[1].class_type;
          const values = e[1].inputs;
          const node = workflow.nodes.find(n => n.id === id);
          let title = node.title;
          if (!title && LiteGraph.registered_node_types[type]) {
            const n = LiteGraph.registered_node_types[type];
            title = n.title;
          }
          return {
            id,
            title,
            type,
            values,
          }
        });

        node.parsedData = {
          filePath,
          dirname,
          fullname,
          filename,
          extname,
          width,
          height,
          nodes: [...nodes, ...notes].sort((a, b) => a.id - b.id),
          prompt,
          workflow,
        }

        return node.parsedData;
      } catch(err) {
        console.error(err);
      }
    }).bind(this);

    this.statics.setMetadata = (async function() {
      const parsedData = await this.statics.getMetadata();
      if (!parsedData) {
        return;
      }

      const {
        filePath,
        fullname,
        dirname,
        filename,
        extname,
        width, 
        height, 
        nodes, 
        workflow, 
        prompt
      } = parsedData;

      const q = this.widgets[0];
      const v = this.widgets[1];

      if (q.value === "PATH") {
        v.value = filePath;
      } else if (q.value === "FILENAME") {
        v.value = fullname;
      } else if (q.value === "DIR") {
        v.value = dirname;
      } else if (q.value === "NAME") {
        v.value = filename;
      } else if (q.value === "EXT") {
        v.value = extname.replace(".", "");
      } else if (q.value === "WIDTH") {
        v.value = width;
      } else if (q.value === "HEIGHT") {
        v.value = height;
      } else {
        let qv = q.value.split(".");
        let nn = qv.slice(0, qv.length - 1).join(".");
        let ni = 0;
        let wn = qv.slice(qv.length - 1).join(".") || ""; // Note.

        // parse index of node
        if (/\[([0-9]+)\]$/.test(nn)) {
          ni = parseInt(/\[([0-9]+)\]$/.exec(nn).pop());
          nn = nn.replace(/\[([0-9]+)\]$/, "");
        }

        let n = findNode(nodes, nn, ni);
        if (n && n.values[wn]) {
          v.value = n.values[wn];
        }
      }
    }).bind(this);

    setTimeout(async () => {
      this.setSize(this.size);
      this.setDirtyCanvas(true, true);

      // First run after creation
      // try {
      //   await updateNodes();
      // } catch(err) {
      //   // console.error(err);
      // }

      this.statics.isInitialized = true;
    }, 128);

    // Prevent update on initialization
    setTimeout(() => {
      this.onConnectionsChange = function() {
        updateNodes();
      }
    }, 1024);
  } catch(err) {
    console.error(err);
  }
}

function initLoader() {

  const setCallback = function(widget) {
    if (!widget) {
      return;
    }
    const orig = widget.callback;
    widget.callback = function () {
      const r = orig ? orig.apply(this, arguments) : undefined;
      updateNodes();
      return r;
    }
  }

  switch(this.comfyClass) {
    // core
    case "LoadImage":
    case "LoadImageMask":
      setCallback(this.widgets?.find(e => e.name === "image"));
      break;
    // ComfyUI-Inspire-Pack
    case "LoadImage //Inspire": 
      setCallback(this.widgets?.find(e => e.name === "image"));
      break;
    // ComfyUI-Crystools
    case "Load image with metadata [Crystools]":
      setCallback(this.widgets?.find(e => e.name === "image"));
      break;
    // WAS Node Suite
    case "Image Load": 
      setCallback(this.widgets?.find(e => e.name === "image_path"));
      break;
    // comfyui-put-image
    case "PutImage": 
      setCallback(this.widgets?.find(e => e.name === "index"));
      setCallback(this.widgets?.find(e => e.name === "dir_path"));
      break; 
  }
}

function getFilePath(node) {
  if (node && node.widgets) {
    let prefix, suffix;
    switch(node.comfyClass) {
      // core
      case "LoadImage":
      case "LoadImageMask": 
        prefix = "ComfyUI/input";
        suffix = node.widgets.find(e => e.name === "image")?.value;
        break;
      // ComfyUI-Inspire-Pack
      case "LoadImage //Inspire": 
        prefix = "ComfyUI/input";
        suffix = node.widgets.find(e => e.name === "image")?.value;
        break;
      // ComfyUI-Crystools
      case "Load image with metadata [Crystools]": 
        prefix = "ComfyUI/input";
        suffix = node.widgets.find(e => e.name === "image")?.value;
        break;
      // WAS Node Suite
      case "Image Load": 
        suffix = node.widgets.find(e => e.name === "image_path")?.value;
        break;
      // comfyui-put-image
      case "PutImage": 
        prefix = node.widgets.find(e => e.name === "dir_path")?.value;
        suffix = node.widgets.find(e => e.name === "filename")?.value;
        break;
    }
    if (prefix && suffix) {
      return `${prefix}/${suffix}`.replace(/[\\/]+/g, "/");
    } else if (suffix) {
      return suffix.replace(/[\\/]+/g, "/");;
    }
  }
}

function matchNode(n, q) {
  if (!n || !q) {
    return false;
  }
  if (n.title && n.title === q) {
    return true;
  }
  if (n.type && n.type === q) {
    return true;
  }
  if (n.id && n.id == q) {
    return true;
  }
  return false;
}

function findNode(nodes, query, index, reverse) {
  if (!index) {
    index = 0;
  }
  let count = 0;
  if (!reverse) {
    for (let i = 0; i < nodes.length; i++) {
      if (matchNode(nodes[i], query)) {
        if (count === index) {
          return nodes[i];
        } else {
          count++;
        }
      }
    }
  } else {
    for (let i = nodes.length - 1; i >= 0; i--) {
      if (matchNode(nodes[i], query)) {
        if (count === index) {
          return nodes[i];
        } else {
          count++;
        }
      }
    }
  }
}

async function updateNodes() {
  // Clear parsedData in image nodes
  for (const node of app.graph._nodes) {
    if (IMAGE_LOADER_TYPES.indexOf(node.comfyClass) > -1) {
      delete node.parsedData;
    }
  }

  // Load all parse nodes
  for (const node of app.graph._nodes) {
    if (PARSER_TYPES.indexOf(node.comfyClass) > -1) {
      try {
        await node.statics.setMetadata();
      } catch(err) {
        console.error(err);
      }
    }
  }
}

app.registerExtension({
	name: `shinich39.ParseImage`,
  setup() {
    // render before start a new queue
    // const origQueuePrompt = app.queuePrompt;
    // app.queuePrompt = async function(number, batchCount) {
    //   await updateNodes();
    //   const r = await origQueuePrompt.apply(this, arguments);
    //   return r;
    // }

    // Append event to last sequence of events
    setTimeout(() => {
      api.addEventListener("promptQueued", async function() {
        await updateNodes();
      });

      console.log("[compfyui-parse-image] Event added.");
    }, 512);
  },
  async afterConfigureGraph(missingNodeTypes) {
    for (const node of app.graph._nodes) {
      if (PARSER_TYPES.indexOf(node.comfyClass) > -1) {
        initParser.apply(node);
      }
    }
	},
  nodeCreated(node) {
    if (PARSER_TYPES.indexOf(node.comfyClass) > -1) {
      initParser.apply(node);
    } else if (IMAGE_LOADER_TYPES.indexOf(node.comfyClass) > -1) {
      // prevent callbacks during initialization
      setTimeout(() => {
        initLoader.apply(node);
      }, 1024);
    }
  },
});